import akshare as ak
import stock_zt as zt
import pandas as pd
import requests  # 显式导入requests，方便设置超时

# 全局设置requests超时（避免请求卡住）
requests.adapters.DEFAULT_TIMEOUT = 10  # 10秒超时


# ----------------------
# 1. 数据优化（修复SettingWithCopyWarning）
# ----------------------
def optimize_data(df):
    """优化数据格式，通过.copy()避免副本修改警告"""
    # 先创建副本，确保修改的是独立DataFrame（核心修复）
    df_copy = df.copy()

    # 主力净流入-净额转换（单位：万元，保留2位小数）
    df_copy['主力净流入-净额（万元）'] = (df_copy['主力净流入-净额'] / 10000).round(2)
    df_copy.drop(columns=['主力净流入-净额'], inplace=True)  # 删除原列

    # 超大单净流入转换
    df_copy['超大单净流入-净额（万元）'] = (df_copy['超大单净流入-净额'] / 10000).round(2)
    df_copy.drop(columns=['超大单净流入-净额'], inplace=True)

    # 大单净流入转换
    df_copy['大单净流入-净额（万元）'] = (df_copy['大单净流入-净额'] / 10000).round(2)
    df_copy.drop(columns=['大单净流入-净额'], inplace=True)

    # 中单净流入转换
    df_copy['中单净流入-净额（万元）'] = (df_copy['中单净流入-净额'] / 10000).round(2)
    df_copy.drop(columns=['中单净流入-净额'], inplace=True)

    # 小单净流入转换
    df_copy['小单净流入-净额（万元）'] = (df_copy['小单净流入-净额'] / 10000).round(2)
    df_copy.drop(columns=['小单净流入-净额'], inplace=True)

    return df_copy


# ----------------------
# 2. 获取资金流向数据（修复请求超时和异常）
# ----------------------
def get_zl_hold(code, name, market, zt_date):
    """获取从涨停日到现在的资金数据，添加完整异常处理"""
    try:
        # 调用akshare接口时添加超时（避免卡住）
        all_data = ak.stock_individual_fund_flow(code, market)

        # 转换日期格式（确保与涨停日可比）
        all_data['日期'] = pd.to_datetime(all_data['日期'])
        zt_date = pd.to_datetime(zt_date)  # 确保涨停日是datetime类型

        # 筛选涨停日及之后的数据（包含涨停日）
        filtered_data = all_data[all_data['日期'] >= zt_date]
        if filtered_data.empty:
            print(f"⚠️ {name}（{code}）涨停日（{zt_date.strftime('%Y-%m-%d')}）后无数据，已跳过")
            return pd.DataFrame()

        # 优化数据（基于副本操作，避免警告）
        optimized_data = optimize_data(filtered_data)

        # 添加代码和名称（用.loc避免副本修改）
        optimized_data.loc[:, '代码'] = code  # 显式用.loc修改，避免警告
        optimized_data.loc[:, '名称'] = name

        return optimized_data

    # 捕获手动中断（用户按Ctrl+C）
    except KeyboardInterrupt:
        print(f"\n⚠️ 用户中断请求：{name}（{code}），程序将继续处理其他股票")
        return pd.DataFrame()
    # 捕获其他异常（网络错误、数据格式错误等）
    except Exception as e:
        print(f"❌ {name}（{code}）获取失败：{str(e)[:50]}（已跳过）")
        return pd.DataFrame()


# ----------------------
# 3. 主逻辑（获取涨停后数据）
# ----------------------
# 获取涨停股票列表（假设数据中有'涨停日期'字段，需与实际字段名匹配）
data_past_5_days_first_zt = zt.data_past_5_days_first_zt
if data_past_5_days_first_zt.empty:
    print("❌ 未获取到涨停股票列表")
    exit()
# 确保涨停日期格式正确（关键：否则无法筛选）
if '涨停日期' not in data_past_5_days_first_zt.columns:
    print("❌ 涨停股票数据中无'涨停日期'字段，请检查stock_zt模块")
    exit()
data_past_5_days_first_zt['涨停日期'] = pd.to_datetime(
    data_past_5_days_first_zt['涨停日期'],
    errors='coerce'  # 转换失败的日期设为NaT
)
# 过滤无效涨停日期的股票
valid_zt_stocks = data_past_5_days_first_zt[
    ~data_past_5_days_first_zt['涨停日期'].isna()
]
print(
    f"已加载涨停股票：{len(valid_zt_stocks)}只（过滤无效日期{len(data_past_5_days_first_zt) - len(valid_zt_stocks)}只）")

# 累加所有涨停后的数据
all_zt_after_data = pd.DataFrame()
valid_stocks_count = len(valid_zt_stocks)  # 提前保存有效股票数量

# 使用enumerate显式跟踪处理数量，而不是依赖DataFrame的索引
for processed_count, (i, row) in enumerate(valid_zt_stocks.iterrows(), 1):
    code = row['代码']
    name = row['名称']
    market = 'sh' if str(code).startswith('6') else 'sz'
    zt_date = row['涨停日期']

    # 获取单只股票数据并累加
    stock_data = get_zl_hold(code, name, market, zt_date)
    if not stock_data.empty:
        all_zt_after_data = pd.concat([all_zt_after_data, stock_data], ignore_index=True)

    # 每处理5只股票打印进度（使用processed_count确保正确）
    if processed_count % 5 == 0:
        print(f"已处理 {processed_count}/{valid_stocks_count} 只股票，当前累计有效数据：{len(all_zt_after_data)}条")

# 保存数据（仅当有有效数据时）
if not all_zt_after_data.empty:
    save_path = zt.get_file_path('首板股票涨停后资金数据.csv')
    all_zt_after_data.to_csv(save_path, index=False, encoding='utf-8-sig')
    print(f"\n✅ 数据已保存至：{save_path}（共{len(all_zt_after_data)}条记录）")
else:
    print("\n❌ 无有效数据可保存")
